from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel

checkpointer = InMemorySaver()

model = init_chat_model(
    "deepseek-chat",
    temperature=0
)

class WeatherResponse(BaseModel):
    conditions: str

def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"It's always sunny in {city}!"

agent = create_react_agent(
    model=model,
    tools=[get_weather],
    prompt="Never answer questions about the weather.",
    checkpointer=checkpointer, # 开启持久化
    response_format=WeatherResponse
)

config = {"configurable": {"thread_id": "1"}}

sf_response = agent.invoke(
    {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
    config
)
ny_response = agent.invoke(
    {"messages": [{"role": "user", "content": "what about new york?"}]},
    config
)

